package com.example.nplboot.xl;

import cn.hutool.core.io.FileUtil;
import com.aliasi.classify.Classification;
import com.aliasi.classify.Classified;
import com.aliasi.classify.TfIdfClassifierTrainer;
import com.aliasi.tokenizer.CharacterTokenizerFactory;
import com.aliasi.tokenizer.TokenFeatureExtractor;
import com.hankcs.hanlp.HanLP;
import com.hankcs.hanlp.seg.common.Term;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

/**
 * 训练方法，直接运行
 * @Classname XlTest
 * @Date 2023/2/10 22:16
 * @Created by 87766867@qq.com
 */
public class XlTest {

    private  static String[] CATEGORIES = {
            "原告被告",
            "原告被告代理人",
            "原告被告代理人代表人",
            "原告被告代表人"

    };
    //训练文件目录
    private static String TEXT_CLASSIFICATION_TRAINING = "/数据目录/training/";
    //测试文件目录
    private static String TEXT_CLASSIFICATION_TESTING = "/数据目录/testing/";
    //模型目录
    private static String MODEL_FILE = "/数据目录/训练好的模型/tfidf_classifier";

    //npl
    private static TfIdfClassifierTrainer<CharSequence> classifier = new TfIdfClassifierTrainer<CharSequence>(new TokenFeatureExtractor(CharacterTokenizerFactory.INSTANCE));

    public static void main(String[] args) throws IOException {
        File directory = new File("");
        try {
            TEXT_CLASSIFICATION_TRAINING=directory.getCanonicalPath()+TEXT_CLASSIFICATION_TRAINING;
            TEXT_CLASSIFICATION_TESTING=directory.getCanonicalPath()+TEXT_CLASSIFICATION_TESTING;
            MODEL_FILE=directory.getCanonicalPath()+MODEL_FILE;
        } catch (IOException e) {
            e.printStackTrace();
        }
        List<List<String>> listdatas=redFile("txt");
        trainModel(true,listdatas);

    }

    /**
     * 训练模型方法
     * @param needStoreModle 是否保存训练模型
     * @throws IOException
     */
    public static void trainModel(boolean needStoreModle,List<List<String>> listdatas) throws IOException {
        for (int i = 0; i < listdatas.size(); i++) {
            for (int j = 0; j <listdatas.get(i).size(); j++) {
                String segWords = "";
                //读取内容后分词
                List<Term> terms = HanLP.segment(listdatas.get(i).get(j));
                for (Term term : terms)
                {
                    segWords += term.word + " ";
                }
                Classification classification = new Classification(CATEGORIES[i]);
                Classified<CharSequence> classified= new Classified<CharSequence>(segWords, classification);
                classifier.handle(classified);
            }
        }
        if (needStoreModle) {
            MODEL_FILE=MODEL_FILE.replaceAll("\\\\","/");
            System.out.println("开始保存分类器到  " + MODEL_FILE);
            ObjectOutputStream os = new ObjectOutputStream(new FileOutputStream(MODEL_FILE));
            classifier.compileTo(os);
            os.close();
            System.out.println("分类器保存完成");
        }
    }

    /**
     * 读取文件方法
     * @param type 文件后缀 .txt
     * @return
     */
    public  static List<List<String>> redFile(String type){
        List<List<String>> listdatas=new ArrayList<List<String>>();
        for (String fl : CATEGORIES){
            List<String> datas=new ArrayList<String>();
            String fileurl=TEXT_CLASSIFICATION_TRAINING+fl+"/".replaceAll("\\\\","/");
            if(!FileUtil.isDirectory(fileurl)){
                String msg = "找不到训练目录" + fileurl;
                System.out.println(msg);
                throw new IllegalArgumentException(msg);
            }
            List<String> list=FileUtil.listFileNames(fileurl);
            for (int i = 0; i < list.size(); i++) {
                if("doc".equals(type.toLowerCase())||"docx".equals(type.toLowerCase())){
                    String str= readWord(fileurl+"/"+list.get(i));
                    datas.add(str);
                }
                if("txt".equals(type.toLowerCase())){
                    String str=FileUtil.readString(fileurl+"/"+list.get(i), StandardCharsets.UTF_8 );
                    datas.add(str);
                }
            }
            listdatas.add(datas);
        }
        return listdatas;

    }


    public static String readWord(String path) {
        String filePath = path;

        String suffixName = filePath.substring(filePath.lastIndexOf("."));//从最后一个.开始截取。截取fileName的后缀名
        try {
            File file = new File(filePath);
            FileInputStream fs = new FileInputStream(file);
            if(suffixName.equalsIgnoreCase(".doc")){//doc
                StringBuilder result2 = new StringBuilder();
                WordExtractor re = new WordExtractor(fs);
                result2.append(re.getText());//获取word中的文本内容
                re.close();
                return result2.toString();

            }else{//docx
                XWPFDocument doc = new XWPFDocument(fs);
                XWPFWordExtractor extractor = new XWPFWordExtractor(doc);
                String text = extractor.getText();//获取word中的文本内容
                extractor.close();
                fs.close();
                return text;
            }
        }catch (Exception e) {
            e.printStackTrace();
        }
        return "";

    }




}
